import Router from '@koa/router';
import axios from 'axios';
import { log } from 'console';
import sse from 'koa-sse-stream';
import { PassThrough } from 'stream';
import { Transform } from 'stream';

class OllamaStreamParser extends Transform {
  constructor() {
    super({ objectMode: true });
    this.buffer = '';
    this.inCodeBlock = false;
    this.language = '';
    this.isThinking = false;
  }

  _transform(chunk, _, callback) {
    try {
      this.buffer += chunk.toString();
      const lines = this.buffer.split('\n');
      this.buffer = lines.pop() || '';

      for (const line of lines) {
        if (!line.trim() || !line.startsWith('data:')) continue;

        const dataStr = line.replace('data:', '').trim();
        console.log('%c [ dataStr ]-27', 'font-size:13px; background:pink; color:#bf2c9f;', dataStr)
        if (dataStr.includes('[DONE]')) return;
        console.log('%c [ !dataStr.includes("[DONE]") ]-29', 'font-size:13px; background:pink; color:#bf2c9f;', !dataStr.includes('[DONE]'))
        console.log('%c [ dataStr ]-27', 'font-size:13px; background:pink; color:#bf2c9f;', dataStr)
        if (!dataStr) continue;

        // 处理可能的SSE格式

        const parsed = JSON.parse(dataStr);
        const isDeepSeek = !!parsed.choices;

        // 统一数据格式处理
        console.log('%c [ isDeepSeek ]-104', 'font-size:13px; background:pink; color:#bf2c9f;', isDeepSeek)
        let data = isDeepSeek ? this.parseDeepSeekData(parsed) : this.parseOllamaData(parsed);
        console.log('%c [ data ]-105', 'font-size:13px; background:pink; color:#bf2c9f;', data)
        const { content, done } = data

        // 处理代码块和思考标记
        const processed = this.processContent(content);

        this.push(JSON.stringify({
          ...processed,
          done: done || false
        }));
      }
    } catch (err) {
      console.error('解析错误', err);
    }
    callback();
  }

  parseDeepSeekData(parsed) {
    return {
      type: this.isThinking ? 'think' : 'text', // 根据状态设置类型
      content: parsed.choices?.[0]?.delta?.content || '',
      done: parsed.choices?.[0]?.finish_reason === 'stop'
    };
  }

  parseOllamaData(parsed) {
    return {
      type: this.isThinking ? 'think' : 'text', // 根据状态设置类型
      content: parsed.response || '',
      done: parsed.done || false
    };
  }

  processContent(content) {

    // 检测思考标记
    if (content.includes('<think>')) {
      this.isThinking = true;
      return { type: 'think', content: '', isCode: false };
    } else if (content.includes('</think>')) {
      this.isThinking = false;
      return { type: 'text', content: '', isCode: false };
    }
    // 处理代码块标记
    if (content.includes('```')) {
      console.log(';检测思考标记检测思考标记2222222', content);
      this.inCodeBlock = !this.inCodeBlock;
      if (this.inCodeBlock) {
        const match = content.match(/```(\w+)?/);
        this.language = match?.[1] || '';
      }
      return { type: 'code_mark', content: '', isCode: this.inCodeBlock };
    }

    // 转义特殊字符（非代码块内容）
    if (!this.inCodeBlock && !content.includes('\n')) {
      content = this.escapeSpecialChars(content);
    }

    return {
      type: this.isThinking ? 'think' : 'text',
      content,
      isCode: this.inCodeBlock
    };
  }

  escapeSpecialChars(str) {
    return str.replace(/[&<>"'`=\/]/g, (match) => ({
      '&': '\\u0026',
      '<': '\\u003c',
      '>': '\\u003e',
      '"': '\\u0022',
      "'": '\\u0027',
      '`': '\\u0060',
      '\r': '\\u000d',
      '\t': '\\u0009'
    }[match] || match));
  }
}

const router = new Router({ prefix: '/api' });

router.get('/chat-stream', async (ctx) => {
  // 1. 手动创建SSE流
  const stream = new PassThrough();
  ctx.set({
    'Content-Type': 'text/event-stream',
    'Cache-Control': 'no-cache',
    'Connection': 'keep-alive'
  });
  ctx.body = stream;

  // 2. 发送初始化事件
  stream.write('event: init\ndata: {}\n\n');
  try {
    // 3. 处理Ollama流
    const ollamaStream = await axios.post('http://127.0.0.1:11434/api/generate', {
      model: 'deepseek-r1:8b',
      prompt: ctx.query.prompt,
      stream: true,
      options: {
        num_ctx: 4096,  // 增大上下文窗口
        num_predict: -1, // 取消长度限制（或设置更大值如2048）
        temperature: 0.7, // 控制随机性
        top_p: 0.9, // 核采样参数
        repeat_penalty: 1.1, // 抑制重复

      }
    }, {
      responseType: 'stream',
      timeout: 120000 // 延长超时到2分钟
    });
    // 4. 创建处理管道
    const parser = new OllamaStreamParser();

    // 错误处理
    const handleError = (err) => {
      console.error('[流错误]', err);
      stream.write(`event: error\ndata: ${JSON.stringify({ error: err.message })}\n\n`);
      stream.end();
    };

    ollamaStream.data.on('error', handleError);
    parser.on('error', handleError);
    stream.on('error', handleError);
    // 4. 管道连接
    ollamaStream.data.pipe(parser)
      .on('data', chunk => {
        stream.write(`data: ${chunk.toString()}\n\n`);
      })
      .on('end', () => {
        stream.write('event: end\ndata: {}\n\n');
        stream.end();
      });
    // 客户端断开连接处理
    ctx.req.on('close', () => {
      ollamaStream.data.destroy();
      console.log('[客户端断开连接]');
    });
  } catch (error) {
    console.error('[请求错误]', error);
    stream.write(`event: error\ndata: ${JSON.stringify({
      error: error.message,
      code: error.response?.status
    })}\n\n`);
    stream.end();
  }

});
// 2. 调用 DeepSeek API
router.get('/chat-stream/api', async (ctx) => {
  // 1. 手动创建SSE流
  const stream = new PassThrough();
  ctx.set({
    'Content-Type': 'text/event-stream',
    'Cache-Control': 'no-cache',
    'Connection': 'keep-alive'
  });
  ctx.body = stream;


  const prompt = ctx.query.prompt;
  console.log('%c [ prompt ]-141', 'font-size:13px; background:pink; color:#bf2c9f;', prompt)
  let response = await axios.post(
    'https://api.deepseek.com/chat/completions',
    {
      model: "deepseek-chat",
      messages: [{ role: "user", content: prompt }],
      stream: true,
      temperature: 0.7, // 推荐添加生成参数
      max_tokens: 2048
    },
    {
      headers: {
        "Authorization": `Bearer ${process.env.DEEPSEEK_KEY}`,
        "Content-Type": "application/json",
        "Accept": "text/event-stream" // 明确接受 SSE 格式
      },
      responseType: "stream",
      timeout: 30000 // 延长超时时间
    }
  );
  const parser = new OllamaStreamParser();
  // 错误处理
  const handleError = (err) => {
    console.error('[流错误]', err);
    stream.write(`event: error\ndata: ${JSON.stringify({ error: err.message })}\n\n`);
    stream.end();
  };
  response.data.on('error', handleError);
  parser.on('error', handleError);
  stream.on('error', handleError);
  response.data.pipe(parser)
    .on('data', chunk => {
      stream.write(`data: ${chunk.toString()}\n\n`);
    })
    .on('end', () => {
      stream.write('event: end\ndata: {}\n\n');
      stream.end();
    });
  // 客户端断开连接处理 
  ctx.req.on('close', () => {
    response.data.destroy();
    console.log('[客户端断开连接]');
  });

  // let buffer = '';
  // response.data.on('data', (chunk) => {

  //   buffer += chunk.toString();

  //   // 处理完整的事件
  //   const events = buffer.split('\n\n');
  //   buffer = events.pop() || ''; // 保留不完整的事件

  //   for (const event of events) {
  //     console.log('%c [ event ]-171', 'font-size:13px; background:pink; color:#bf2c9f;', event)
  //     try {
  //       const lines = event.split('\n').filter(line => line.trim());
  //       for (const line of lines) {
  //         if (line.startsWith('data:')) {
  //           const data = JSON.parse(line.substring(5).trim());
  //           if (data.choices?.[0]?.delta?.content) {
  //             ctx.sse.send(data.choices[0].delta.content);
  //           }
  //           // 处理完成事件
  //           if (data.choices?.[0]?.finish_reason) {
  //             ctx.sse.send('[DONE]');
  //           }
  //         }
  //       }
  //     } catch (e) {
  //       console.error('Parse error:', e.message);
  //     }
  //   }
  // });


})


export default router;